#!/usr/bin/env python
# encoding: utf-8
"""
main.py

Created by Brett Harrison and Daniel Suo
"""

import sys
import getopt

from ppm import read
from image_util import *
from algorithm import *
from vec import sumsqrs
from time import time
import numpy as np
import scikits.ann as ANN

help_message = '''
The help message goes here.
'''


class Usage(Exception):
	def __init__(self, msg):
		self.msg = msg


def main(argv=None):
	if argv is None:
		argv = sys.argv
	try:
		try:
			opts, args = getopt.getopt(argv[1:], "ho:v", ["help", "output="])
		except getopt.error, msg:
			raise Usage(msg)
	
		# option processing
		for option, value in opts:
			if option == "-v":
				verbose = True
			if option in ("-h", "--help"):
				raise Usage(help_message)
			if option in ("-o", "--output"):
				output = value
	
	except Usage, err:
		print >> sys.stderr, sys.argv[0].split("/")[-1] + ": " + str(err.msg)
		print >> sys.stderr, "\t for help use --help"
		return 2
	
def inpainting(filename="../asst2io/inpainting_minor_clean.ppm", patchSize=9):
	GREEN = [[(0,255,0) for i in range(patchSize)] for j in range(patchSize)]
	print "Reading image..."
	#image = read(filename)
	#image = Image.open("../asst2io/test.ppm")
	#image = Image.open("../asst2io/lauren2.pbm")
	#image = Image.open("../asst2io/inpainting_minor_shrunk.pbm")
	image = Image.open("../asst2io/inpainting_major_shrunk.pbm")
	
	width,height = image.size
	print "Converting to CIELab..."
	rgbMatrix = imageToList(image)
	labMatrix = [[RgbToLab(rgb) for rgb in row] for row in rgbMatrix]
	
	print "Precomputing patches..."
	patches, indexMap = initPatches(labMatrix, patchSize)
	
	patches = np.array(patches)
	
	print "Initializing pixel confidence values..."
	confidenceMatrix = initConfidenceMatrix(rgbMatrix)
	
	print "Precomputing gradient values..."
	gradientMatrix = []
	for i in range(width):
		row = []
		for j in range(height):
			if i == 0 or i == width-1 or j == 0 or j == height-1:
				row.append((0.0,0.0,0.0))
			elif not isBlack(rgbMatrix,i,j) and filter(lambda (x,y): isBlack(rgbMatrix,x,y), getNeighbors(i,j)) == []:
				row.append(imageGradient(labMatrix,i,j))
			else:
				row.append((0.0,0.0,0.0))
		gradientMatrix.append(row)
	
	#for ITER in range(3):
	ITER = 1
	while True:
		print "### ITERATION %d ###" % ITER
		print "Computing fill front..."
		fillFront = filter(lambda (i,j): isBoundary(rgbMatrix,i,j), [(x,y) for x in range(1,width-1) for y in range(1,height-1)])
		
		if fillFront == []:
			break
		
		binaryMatrix = [[c if c == (0,0,0) else (255,255,255) for c in row] for row in rgbMatrix]
		print "%d boundary pixels left." % (len(fillFront))
		print "Blurring image..."
		blurImage = blur(listToImage(binaryMatrix))
		
		maxPriority = -1
		maxPatchX = maxPatchY = None
	
		print "Finding border patch with max priority..."
		for (i,j) in fillFront:
			priority = dataTerm(blurImage,gradientMatrix,i,j,patchSize)*confidenceTerm(confidenceMatrix,i,j,patchSize)
			if priority > maxPriority:
				maxPriority = priority
				maxPatchX, maxPatchY = (i,j)
				
		print "Finding exemplar patch..."
		t = time()
		maxPatch = np.array(getLabPatch(labMatrix, maxPatchX, maxPatchY, patchSize))
		nonblack = np.array(map(lambda x: map(lambda y: y[0], x), maxPatch)).nonzero()
		maxPatch = maxPatch[nonblack]

		print "Removed columns in %f seconds." % (time()-t)
		t = time()

		maxPatchArray = np.array([maxPatch for p in range(len(patches))])
		newPatches = map(lambda x: x[nonblack], patches)

		diff = maxPatchArray - newPatches
		diff = diff * diff
		exemplarX, exemplarY = indexMap[np.array(map(lambda x: x.sum(), diff)).argmin()]
		exemplar = getPatch(labMatrix, exemplarX, exemplarY, patchSize)
		print "Found exemplar in %f seconds." % (time()-t)
		
		print "Copying image data from exemplar..."
		newPixels = fillPatch(labMatrix,exemplar,maxPatchX,maxPatchY)
		
		print "Updating confidence..."
		newConfidence = confidenceTerm(confidenceMatrix,maxPatchX,maxPatchY,patchSize)
		for (i,j) in newPixels:
			confidenceMatrix[i][j] = newConfidence
			
		print "Reconverting to RGB..."
		rgbMatrix = [[LabToRgb(lab) for lab in row] for row in labMatrix]
		image = listToImage(rgbMatrix)
		ITER += 1
	image.save("inpainting_major_shrunk_result.png")
	image.show(command="xv")

if __name__ == "__main__":
	sys.exit(inpainting())
	#sys.exit(main())	